import os
import time
import torch
import numpy as np
from sklearn.metrics import f1_score, accuracy_score
from torch.utils.data import DataLoader
from tqdm import tqdm
from skimage.filters import threshold_otsu
import matplotlib.pyplot as plt

# 从原始代码导入必要组件
from main import (
    config,
    TuSimpleDataset,
    LaneDetectionNet,
    transforms
)

# 定义测试专用预处理（无数据增强）
test_transform = transforms.Compose([
    transforms.Resize(config['input_size']),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])


def visualize_predictions(model, dataloader, device, num_samples=3):
    """可视化预测结果以辅助调试"""
    model.eval()
    images, labels = next(iter(dataloader))

    plt.figure(figsize=(15, 5 * num_samples))
    for i in range(num_samples):
        img = images[i].unsqueeze(0).to(device)
        label = labels[i].cpu().numpy().squeeze()

        with torch.no_grad():
            pred = torch.sigmoid(model(img)).cpu().numpy().squeeze()

        # 自动计算阈值
        try:
            thresh = threshold_otsu(pred)
        except:
            thresh = 0.5
        binary = (pred > thresh).astype(np.uint8)

        # 可视化
        plt.subplot(num_samples, 3, i * 3 + 1)
        plt.imshow(images[i].permute(1, 2, 0).cpu().numpy() * 0.5 + 0.5)
        plt.title("Input Image")

        plt.subplot(num_samples, 3, i * 3 + 2)
        plt.imshow(label, cmap='gray')
        plt.title("Ground Truth")

        plt.subplot(num_samples, 3, i * 3 + 3)
        plt.imshow(binary, cmap='gray')
        plt.title(f"Prediction (Thresh={thresh:.2f})")

    plt.tight_layout()
    plt.show()


def evaluate_model():
    test_config = {
        'batch_size': 8,
        'device': config['device'],
        'warmup_runs': 10,
        'test_runs': 100
    }

    # 加载测试集
    _, test_files = TuSimpleDataset.split_dataset(config['data_root'], test_size=0.2)
    test_set = TuSimpleDataset(config['data_root'], test_files, transform=test_transform)
    test_loader = DataLoader(test_set,
                             batch_size=test_config['batch_size'],
                             num_workers=config['num_workers'],
                             pin_memory=True)

    # 加载模型
    model = LaneDetectionNet().to(test_config['device'])
    model.load_state_dict(torch.load('best_model.pth'))
    model.eval()

    # 先进行可视化检查
    visualize_predictions(model, test_loader, test_config['device'])

    def calculate_metrics(y_true, y_pred):
        y_true_np = y_true.flatten().cpu().numpy().astype(np.uint8)  # 确保转为整型
        y_pred_np = y_pred.flatten().cpu().numpy()

        # 处理全零预测
        if len(np.unique(y_pred_np)) == 1:
            return 0, 0, 0

        # 动态阈值
        try:
            thresh = threshold_otsu(y_pred_np)
        except:
            thresh = 0.5
        y_pred_bin = (y_pred_np > thresh).astype(np.uint8)  # 确保转为整型

        # 计算混淆矩阵（更稳健的写法）
        tp = np.sum((y_true_np == 1) & (y_pred_bin == 1))
        tn = np.sum((y_true_np == 0) & (y_pred_bin == 0))
        fp = np.sum((y_true_np == 0) & (y_pred_bin == 1))
        fn = np.sum((y_true_np == 1) & (y_pred_bin == 0))

        # 计算指标
        accuracy = accuracy_score(y_true_np, y_pred_bin)
        f1 = f1_score(y_true_np, y_pred_bin, zero_division=0)
        fpr = fp / (fp + tn + 1e-9)

        return accuracy, f1, fpr

    # 指标计算
    total_metrics = {'accuracy': 0, 'f1': 0, 'fpr': 0}
    valid_batches = 0
    positive_pixels = 0
    total_pixels = 0

    with torch.no_grad():
        for images, labels in tqdm(test_loader, desc="Evaluating"):
            images = images.to(test_config['device'])
            labels = labels.to(test_config['device'])

            outputs = model(images)
            preds = torch.sigmoid(outputs)

            # 统计正样本比例
            positive_pixels += torch.sum(labels).item()
            total_pixels += labels.numel()

            # 计算指标
            acc, f1, fpr = calculate_metrics(labels, preds)
            if not np.isnan(acc):  # 跳过无效批次
                total_metrics['accuracy'] += acc
                total_metrics['f1'] += f1
                total_metrics['fpr'] += fpr
                valid_batches += 1

    # 平均指标
    avg_metrics = {k: v / valid_batches for k, v in total_metrics.items()}
    positive_ratio = positive_pixels / total_pixels

    # 推理速度测试
    dummy_input = torch.randn(test_config['batch_size'], 3,
                              *config['input_size']).to(test_config['device'])

    # 预热
    for _ in range(test_config['warmup_runs']):
        _ = model(dummy_input)

    # 正式测试
    torch.cuda.synchronize()
    start_time = time.time()
    for _ in range(test_config['test_runs']):
        _ = model(dummy_input)
    torch.cuda.synchronize()
    fps = (test_config['test_runs'] * test_config['batch_size']) / (time.time() - start_time)

    # 输出结果
    print("\n================= 诊断信息 =================")
    print(f"有效测试批次: {valid_batches}/{len(test_loader)}")
    print(f"正样本像素占比: {positive_ratio:.4%}")
    print(f"正样本比例: {np.mean(y_true_np):.4%}")
    print(f"预测正样本比例: {np.mean(y_pred_bin):.4%}")
    print("\n================ 评估结果 ================")
    print(f"Accuracy: {avg_metrics['accuracy']:.4f}")
    print(f"F1 Score: {avg_metrics['f1']:.4f}")
    print(f"FP Rate:  {avg_metrics['fpr']:.4f}")
    print(f"推理速度: {fps:.2f} FPS")
    print("===========================================")


if __name__ == "__main__":
    evaluate_model()